import logging
from pathlib import Path
from typing import AsyncIterator, List, Optional
from pydantic import BaseModel

from gchat_processing.gch_chain_wrapper import GigaChainWrapper
from .base import BaseBlock, ProgressChunk
from .sources import SourcesItem
from .topics import TopicsItem
from ...config import PaperAIConfig


class ReviewStartRequest(BaseModel):
    state_id: Optional[str] = None
    append: bool = False


class ReviewItem(BaseModel):
    text: str


class ReviewUpdateRequest(BaseModel):
    review: str
    state_id: Optional[str] = None


class ReviewBlock(BaseBlock):
    config: PaperAIConfig
    """Literature review block.

    """

    def __init__(self, idea: str, sources: List[SourcesItem], topics: List[TopicsItem],
                 review: str,
                 request: ReviewStartRequest):
        self.idea = idea
        self.sources = sources
        self.topics = topics
        self.review = review
        self.request = request
        self.gchain_wrapper = GigaChainWrapper(
            sql_cache_path=Path('gchat_processing/sqlite_cache.db'),
            giga_token=self.config.gigachat_token,
            giga_scope='GIGACHAT_API_CORP'
        )
        self.logger = logging.getLogger(__name__)

    @classmethod
    def from_state(cls, data: dict, request: ReviewStartRequest) -> 'ReviewBlock':
        return ReviewBlock(
            idea=data.get('idea', ''),
            sources=[
                SourcesItem.model_validate(item)
                for item in data.get('sources', [])
            ],
            topics=[
                TopicsItem.model_validate(item)
                for item in data.get('topics', [])
            ],
            review=data.get('review', ''),
            request=request,
        )

    async def run(self) -> AsyncIterator[ProgressChunk]:
        results: List[ReviewItem] = []
        if self.request.append:
            results = [self.review]

        self.logger.info('received related texts')
        summary_list = [
            dict(title=paper.title, summary=paper.summary, bibtex=paper.bibtex)
            for paper in self.sources
            if not paper.disabled
        ]
        rel_texts: List[str] = [
            await self.gchain_wrapper.get_rel_works(
                description=self.idea,
                summary_list=summary_list
            )
        ]
        self.logger.info('related texts')
        self.logger.info(rel_texts[0])
        for i, rel_text in enumerate(rel_texts):
            self.logger.info(rel_text)
            new_item = ReviewItem.model_validate({
                'text': rel_text
            })
            results.append(new_item)
            yield ProgressChunk(
                done=False,
                total_steps=len(rel_texts),
                current_step=i + 1,
                data={
                    'topics': [new_item],
                }
            )
        self.review = ''.join([review.text for review in results])
        yield ProgressChunk(
            done=True,
            total_steps=len(rel_texts),
            current_step=len(rel_texts),
            data={
                'review': self.review,
            }
        )

    def to_state(self, data: dict):
        data['review'] = self.review
        return
